{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c6ad424",
   "metadata": {},
   "source": [
    "# 04 Debugging\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/debugging.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d5cf012e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c289b68",
   "metadata": {},
   "source": [
    "## Print\n",
    "\n",
    "`jax.debug.print()` can be used for traced array values within `jax.jit`, `jax.vmap`, and others. Python `print()` will not work here because it will only execute at trace-time (and will print out the abstract tracer object types instead).\n",
    "\n",
    "Python `print()` should be used otherwise for static values such as dtypes and array shapes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "5dbf4e98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x=2.0\n",
      "y=0.9092974066734314\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(0.9092974, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f(x):\n",
    "    jax.debug.print(\"x={x}\", x=x)\n",
    "    y = jnp.sin(x)\n",
    "    jax.debug.print(\"y={y}\", y=y)\n",
    "    return y\n",
    "\n",
    "\n",
    "f(2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8302725",
   "metadata": {},
   "source": [
    "Similarly, with `jax.vmap` only `jax.debug.print()` will print the values being mapped over.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "4162bae1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x=0\n",
      "x=1\n",
      "x=2\n",
      "y=0.0\n",
      "y=0.8414709568023682\n",
      "y=0.9092974066734314\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([0.        , 0.84147096, 0.9092974 ], dtype=float32)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.vmap\n",
    "def f(x):\n",
    "    jax.debug.print(\"x={x}\", x=x)\n",
    "    y = jnp.sin(x)\n",
    "    jax.debug.print(\"y={y}\", y=y)\n",
    "    return y\n",
    "\n",
    "\n",
    "xs = jnp.arange(3)\n",
    "f(xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f14be950",
   "metadata": {},
   "source": [
    "Note that all `x` prints occur first, followed by all `y` prints.\n",
    "\n",
    "Instead, if we want sequential prints, we can use `jax.lax.map()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "9a029001",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y=0.0\n",
      "x=0\n",
      "y=0.8414709568023682\n",
      "x=1\n",
      "y=0.9092974066734314\n",
      "x=2\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([0.        , 0.84147096, 0.9092974 ], dtype=float32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax import lax\n",
    "\n",
    "\n",
    "def f(x):\n",
    "    jax.debug.print(\"x={x}\", x=x)\n",
    "    y = jnp.sin(x)\n",
    "    jax.debug.print(\"y={y}\", y=y)\n",
    "    return y\n",
    "\n",
    "\n",
    "xs = jnp.arange(3)\n",
    "lax.map(f, xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ca50d72",
   "metadata": {},
   "source": [
    "Notice now that the ordering of the prints is not sequential (i.e., `x` is not printed before `y`). This is a byproduct of JAX deeming that `y` depends on `x`.\n",
    "\n",
    "If we need the original order, we can add `ordered=True`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "5bd9750b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x=0\n",
      "y=0.0\n",
      "x=1\n",
      "y=0.8414709568023682\n",
      "x=2\n",
      "y=0.9092974066734314\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([0.        , 0.84147096, 0.9092974 ], dtype=float32)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def f(x):\n",
    "    jax.debug.print(\"x={x}\", x=x, ordered=True)\n",
    "    y = jnp.sin(x)\n",
    "    jax.debug.print(\"y={y}\", y=y, ordered=True)\n",
    "    return y\n",
    "\n",
    "\n",
    "xs = jnp.arange(3)\n",
    "lax.map(f, xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ceae176",
   "metadata": {},
   "source": [
    "## Breakpoint\n",
    "\n",
    "Use `jax.debug.breakpoint()` to get PDB-style debugging that will pause program execution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "620371c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering jdb:\n",
      "Entering jdb:\n",
      "Array(2., dtype=float32)\n",
      "Array(0.9092974, dtype=float32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Array(0.9092974, dtype=float32, weak_type=True),\n",
       " Array(-0.41614684, dtype=float32, weak_type=True))"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jax.jit\n",
    "def f(x):\n",
    "    jax.debug.breakpoint()\n",
    "    y, z = jnp.sin(x), jnp.cos(x)\n",
    "    jax.debug.breakpoint()\n",
    "    return y, z\n",
    "\n",
    "\n",
    "f(2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b2967e",
   "metadata": {},
   "source": [
    "For value-dependent breakpoints, we can use `jax.lax.select()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "adba5ec5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entering jdb:\n",
      "Array(2., dtype=float32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(0.9092974, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def break_if_nonzero(x):\n",
    "    def true_fun(x):\n",
    "        pass\n",
    "\n",
    "    def false_fun(x):\n",
    "        jax.debug.breakpoint()\n",
    "\n",
    "    jax.lax.cond(x == 0, true_fun, false_fun, x)\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def f(x):\n",
    "    break_if_nonzero(x)\n",
    "    return jnp.sin(x)\n",
    "\n",
    "\n",
    "f(2.0)  # Non-zero input"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6453d0f3",
   "metadata": {},
   "source": [
    "## Callbacks\n",
    "\n",
    "If we need more control, we can use `jax.debug.callback()` which gives greater control over the host-side logic via a Python callback. It is compatible with every transformation (`jax.vmap()`, `jax.jit()`, etc.):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "17db9679",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Logging value: 2.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array(-0.41614684, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import logging\n",
    "\n",
    "\n",
    "def log_x(x):\n",
    "    logging.warning(f\"Logging value: {x}\")\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def f(x):\n",
    "    jax.debug.callback(log_x, x)\n",
    "    return jnp.cos(x)\n",
    "\n",
    "\n",
    "f(2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "add43905",
   "metadata": {},
   "source": [
    "Here’s an example in `jax.vmap()` transformation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "456fb280",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Logging value: 0\n",
      "WARNING:root:Logging value: 1\n",
      "WARNING:root:Logging value: 2\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([ 1.        ,  0.5403023 , -0.41614684], dtype=float32)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def log_x(x):\n",
    "    logging.warning(f\"Logging value: {x}\")\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def f(x):\n",
    "    jax.debug.callback(log_x, x)\n",
    "    return jnp.cos(x)\n",
    "\n",
    "\n",
    "batch_f = jax.vmap(f, in_axes=0)\n",
    "\n",
    "xs = jnp.stack(jnp.arange(3))\n",
    "batch_f(xs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
